import numpy as np
import matplotlib.pyplot as plt

def simple_resample(particles, weights, gamma=0.75):
    N = len(particles)
    indexes = np.argsort(weights)[:round(N*0.75)][np.random.permutation(N)%round(N*0.75)]
    # resample according to indexes
    weights = gamma*1/N + (1-gamma)*weights[indexes]
    return indexes, weights

class ParticularFiltering:
    def __init__(self, x=[0,0,0,0,0], n=100,
                 pnames = ["铰链", "拉绳", "密封条", "锁扣", "弹簧"],
                 noise=[1e-20,1e-20,1e-20,1e-20,1e-20],
                 alpha=0.1, beta=0.2, gamma=0.75,
                 thres=1):
        self._n = int(n)
        self._pnames = pnames
        self._pre_trend = np.array(noise, dtype="float32")/20
        self._linear = np.random.randn(self._n * len(x)).reshape((self._n, len(x))) * np.array(noise, dtype="float32")
        self._x_Ps = np.array(x, dtype="float32") + self._linear
        self._alpha = alpha
        self._beta = beta
        self._gamma = gamma
        self._thres = thres
        noise = np.std(self._x_Ps, axis=0)
        self._weight = np.exp(- self._linear ** 2 / (2 * noise)) / (2 * np.pi * noise) ** 0.5
        self._weight /= np.sum(self._weight)

    @property
    def x(self):
        #_linear = np.sum(self._weight*self._linear, axis=0)
        return np.sum(self._weight*self._x_Ps, axis=0)

    def predict(self):
        _linear = np.sum(self._weight*self._linear, axis=0)
        _x = np.sum(self._weight*self._x_Ps, axis=0)
        time_ =  np.floor(np.maximum(0, (self._thres - _x)/np.maximum(_linear, self._pre_trend)))
        return {p: {"healthState": round((1-x)*100, 3),
                    "anomTime": f"{t: .3e}次" if t > 1000 else f"{t}次"}
                for p, t, x in zip(self._pnames, time_, _x)}

    def step(self, x):
        _x = self._x_Ps + self._linear
        noise = np.maximum(np.sum(self._weight[:,None]*(_x - np.sum(self._weight*self._x_Ps, axis=0)) ** 2, axis=0) ** 0.5, self._pre_trend)
        self._weight = np.exp(- (_x - np.sum(self._weight*self._x_Ps, axis=0)) ** 2 / (2 * noise)) / (2 * np.pi * noise) ** 0.5
        self._weight /= np.sum(self._weight,axis=0)
        # 重采样
        self._linear = self._alpha * self._linear + (1 - self._alpha) * (np.array(x, dtype="float32") - _x)
        self._x_Ps = self._beta*(self._x_Ps + self._linear) + (1 - self._beta)*_x
        for i in range(len(self._pnames)):
            inds, self._weight[:,i] = simple_resample(np.arange(len(self._x_Ps)), self._weight[:,i], gamma=self._gamma)
            self._x_Ps[:,i] = self._x_Ps[inds,i]
            self._linear[:,i] = self._linear[inds,i]
        self._weight /= np.sum(self._weight,axis=0)

if __name__ == "__main__":
    pf = ParticularFiltering()
    x = np.array([0,0,0,0,0], dtype="float32")
    x_out = []
    x_est_out = []
    T = np.arange(75)
    for t in T:
        x = np.maximum(0, np.minimum(1, x + np.random.uniform(-1/75, 2/75, 5)))
        pf.step(x)
        # 保存数据
        x_out.append(x)
        x_est_out.append(pf.x)
    print(pf.predict())
    import matplotlib.pyplot as plt
    
    x_out = np.array(x_out)
    x_est_out = np.array(x_est_out)
    print(x_est_out.shape)

    # 显示粒子轨迹、真实值、估计值
    plt.figure()
    for i in range(1):
        plt.plot(T, x_out[:,i], color='green', linewidth=2., label=f"P{i+1} - true value")
        plt.plot(T, x_est_out[:,i], color='red', linewidth=2., label=f"P{i+1} - estimate value")
    plt.legend()
    plt.show()
